{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pygame\n",
    "import numpy as np\n",
    "import random\n",
    "\n",
    "# 游戏环境\n",
    "class SnakeGame:\n",
    "    def __init__(self, width=10, height=10):\n",
    "        self.width = width\n",
    "        self.height = height\n",
    "        self.reset()\n",
    "\n",
    "    def reset(self):\n",
    "        self.snake = [(self.width // 2, self.height // 2)]\n",
    "        self.direction = (0, 0)\n",
    "        self.food = self._place_food()\n",
    "        self.score = 0\n",
    "        self.game_over = False\n",
    "        return self._get_state()\n",
    "\n",
    "    def _place_food(self):\n",
    "        while True:\n",
    "            food = (random.randint(0, self.width - 1), random.randint(0, self.height - 1))\n",
    "            if food not in self.snake:\n",
    "                return food\n",
    "\n",
    "    def _get_state(self):\n",
    "        state = np.zeros((self.width, self.height), dtype=np.float32)\n",
    "        for x, y in self.snake:\n",
    "            state[x, y] = 1.0\n",
    "        state[self.food[0], self.food[1]] = 0.5\n",
    "        return state\n",
    "\n",
    "def step(self, action):\n",
    "    if self.game_over:\n",
    "        return self._get_state(), self.score, self.game_over\n",
    "\n",
    "    # 更新方向\n",
    "    if action == 0:  # 上\n",
    "        new_dir = (-1, 0)\n",
    "    elif action == 1:  # 下\n",
    "        new_dir = (1, 0)\n",
    "    elif action == 2:  # 左\n",
    "        new_dir = (0, -1)\n",
    "    elif action == 3:  # 右\n",
    "        new_dir = (0, 1)\n",
    "    else:\n",
    "        new_dir = self.direction\n",
    "\n",
    "    # 检查方向是否相反\n",
    "    if (new_dir[0] == -self.direction[0] and new_dir[1] == -self.direction[1]):\n",
    "        new_dir = self.direction\n",
    "\n",
    "    self.direction = new_dir\n",
    "\n",
    "    # 移动蛇\n",
    "    new_head = (self.snake[0][0] + self.direction[0], self.snake[0][1] + self.direction[1])\n",
    "\n",
    "    # 检查碰撞\n",
    "    if (new_head[0] < 0 or new_head[0] >= self.width or\n",
    "        new_head[1] < 0 or new_head[1] >= self.height or\n",
    "        new_head in self.snake):\n",
    "        self.game_over = True\n",
    "        return self._get_state(), -1, self.game_over  # 碰撞惩罚\n",
    "\n",
    "    self.snake.insert(0, new_head)\n",
    "\n",
    "    # 检查是否吃到食物\n",
    "    if new_head == self.food:\n",
    "        self.score += 1\n",
    "        self.food = self._place_food()\n",
    "        reward = 1  # 吃到食物的奖励\n",
    "    else:\n",
    "        self.snake.pop()\n",
    "        reward = 0.01  # 每存活一步的奖励\n",
    "\n",
    "    return self._get_state(), reward, self.game_over"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "\n",
    "# 神经网络模型\n",
    "class QNetwork(nn.Module):\n",
    "    def __init__(self, input_size, hidden_size, output_size):\n",
    "        super(QNetwork, self).__init__()\n",
    "        self.fc1 = nn.Linear(input_size, hidden_size)\n",
    "        self.fc2 = nn.Linear(hidden_size, hidden_size)\n",
    "        self.fc3 = nn.Linear(hidden_size, output_size)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = torch.relu(self.fc1(x))\n",
    "        x = torch.relu(self.fc2(x))\n",
    "        x = self.fc3(x)\n",
    "        return x\n",
    "\n",
    "# Q-learning 算法\n",
    "class QLearning:\n",
    "    def __init__(self, input_size, hidden_size, output_size, lr=0.001, gamma=0.99):\n",
    "        self.q_network = QNetwork(input_size, hidden_size, output_size)\n",
    "        self.optimizer = optim.Adam(self.q_network.parameters(), lr=lr)\n",
    "        self.gamma = gamma\n",
    "\n",
    "    def get_action(self, state, epsilon=0.1):\n",
    "        if random.random() < epsilon:\n",
    "            return random.randint(0, 3)\n",
    "        else:\n",
    "            state = torch.FloatTensor(state.flatten())\n",
    "            q_values = self.q_network(state)\n",
    "            return torch.argmax(q_values).item()\n",
    "\n",
    "    def train(self, state, action, reward, next_state, done):\n",
    "        state = torch.FloatTensor(state.flatten())\n",
    "        next_state = torch.FloatTensor(next_state.flatten())\n",
    "        q_values = self.q_network(state)\n",
    "        next_q_values = self.q_network(next_state)\n",
    "\n",
    "        target = reward + (1 - done) * self.gamma * torch.max(next_q_values)\n",
    "        loss = nn.MSELoss()(q_values[action], target)\n",
    "\n",
    "        self.optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        self.optimizer.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "ename": "AttributeError",
     "evalue": "'SnakeGame' object has no attribute 'step'",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[27], line 21\u001b[0m\n\u001b[0;32m     19\u001b[0m \u001b[38;5;28;01mwhile\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m done:\n\u001b[0;32m     20\u001b[0m     action \u001b[38;5;241m=\u001b[39m q_learning\u001b[38;5;241m.\u001b[39mget_action(state, epsilon)\n\u001b[1;32m---> 21\u001b[0m     next_state, reward, done \u001b[38;5;241m=\u001b[39m \u001b[43menv\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mstep\u001b[49m(action)\n\u001b[0;32m     22\u001b[0m     q_learning\u001b[38;5;241m.\u001b[39mtrain(state, action, reward, next_state, done)\n\u001b[0;32m     23\u001b[0m     state \u001b[38;5;241m=\u001b[39m next_state\n",
      "\u001b[1;31mAttributeError\u001b[0m: 'SnakeGame' object has no attribute 'step'"
     ]
    }
   ],
   "source": [
    "# 训练参数\n",
    "num_episodes = 1000\n",
    "epsilon = 1.0\n",
    "epsilon_decay = 0.995\n",
    "min_epsilon = 0.01\n",
    "\n",
    "# 初始化环境和Q-learning\n",
    "env = SnakeGame()\n",
    "input_size = env.width * env.height\n",
    "output_size = 4\n",
    "q_learning = QLearning(input_size, 128, output_size)\n",
    "\n",
    "# 训练循环\n",
    "for episode in range(num_episodes):\n",
    "    state = env.reset()\n",
    "    total_reward = 0\n",
    "    done = False\n",
    "\n",
    "    while not done:\n",
    "        action = q_learning.get_action(state, epsilon)\n",
    "        next_state, reward, done = env.step(action)\n",
    "        q_learning.train(state, action, reward, next_state, done)\n",
    "        state = next_state\n",
    "        total_reward += reward\n",
    "\n",
    "    epsilon = max(min_epsilon, epsilon * epsilon_decay)\n",
    "    print(f\"Episode: {episode}, Total Reward: {total_reward}, Epsilon: {epsilon}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "def render(env):\n",
    "    pygame.init()\n",
    "    cell_size = 40\n",
    "    screen = pygame.display.set_mode((env.width * cell_size, env.height * cell_size))\n",
    "    clock = pygame.time.Clock()\n",
    "\n",
    "    while True:\n",
    "        for event in pygame.event.get():\n",
    "            if event.type == pygame.QUIT:\n",
    "                pygame.quit()\n",
    "                return\n",
    "\n",
    "        screen.fill((0, 0, 0))\n",
    "        for x, y in env.snake:\n",
    "            pygame.draw.rect(screen, (0, 255, 0), (y * cell_size, x * cell_size, cell_size, cell_size))\n",
    "        pygame.draw.rect(screen, (255, 0, 0), (env.food[1] * cell_size, env.food[0] * cell_size, cell_size, cell_size))\n",
    "        pygame.display.flip()\n",
    "        clock.tick(10)\n",
    "\n",
    "# 在训练结束后可视化游戏\n",
    "render(env)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
